import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    """
    LoRA (Low-Rank Adaptation) Linear Layer.
    实现了LoRA (低秩适配) 的线性层，通过低秩矩阵对原始权重矩阵进行适配，从而提高模型训练效率并降低计算复杂度。
    """

    def __init__(self, in_features, out_features, r=8, alpha=32, dropout=0.1, bias=False):
        """
        初始化LoRALinear层。

        参数：
            in_features (int): 输入特征的维度（即输入向量的大小）。
            out_features (int): 输出特征的维度（即输出向量的大小）。
            r (int, optional): 低秩矩阵的秩。默认为8。
            alpha (int, optional): 缩放因子，控制LoRA更新的大小。默认为32。
            dropout (float, optional): dropout率，用于丢弃某些输入特征。默认为0.1。
            bias (bool, optional): 是否使用偏置项。默认为False。
        """
        super(LoRALinear, self).__init__()

        # 输入特征维度和输出特征维度
        self.in_features = in_features
        self.out_features = out_features

        # 低秩矩阵的秩和缩放因子
        self.r = r
        self.alpha = alpha
        self.scaling = self.alpha / self.r  # 计算缩放因子

        # 定义权重参数，不参与梯度更新
        self.weight = nn.Parameter(torch.randn(out_features, in_features))  # 原始权重
        self.weight.requires_grad = False  # 禁止原始权重参与梯度计算

        # 定义LoRA的低秩矩阵 A 和 B
        self.A = nn.Parameter(torch.randn(r, in_features) * 0.01)  # 低秩矩阵A，形状 (r, in_features)
        self.B = nn.Parameter(torch.randn(out_features, r) * 0.01)  # 低秩矩阵B，形状 (out_features, r)

        # Dropout层，用于正则化，避免过拟合
        self.dropout = nn.Dropout(dropout)

        # 定义偏置项（如果需要）
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))  # 偏置初始化为0
        else:
            self.bias = None

    def forward(self, x):
        """
        前向传播，计算输出。

        参数：
            x (torch.Tensor): 输入张量，形状为 (batch_size, in_features)。

        返回：
            torch.Tensor: 输出张量，形状为 (batch_size, out_features)。
        """
        # 计算原始线性变换： x * weight.T
        result = torch.matmul(x, self.weight.T)

        # 计算LoRA的适配更新： x * A.T * B.T，并应用缩放因子
        lora_update = torch.matmul(self.dropout(x), self.A.T)  # x * A.T
        lora_update = torch.matmul(lora_update, self.B.T)  # (x * A.T) * B.T
        lora_update = lora_update * self.scaling  # 应用缩放因子

        # 将LoRA的适配更新加到原始线性变换结果上
        result += lora_update

        # 如果有偏置项，加入偏置
        if self.bias is not None:
            result += self.bias

        return result

